"""
author:dlr123
date:2021年09月28日
"""

import numpy as np


class TreeNode:
    @staticmethod
    def softmax(x):
        if x.ndim == 2:
            x = x - x.max(axis=1, keepdims=True)
            x = np.exp(x)
            x /= x.sum(axis=1, keepdims=True)
        elif x.ndim == 1:
            x = x - np.max(x)
            x = np.exp(x)/np.sum(np.exp(x))
        return x

    def __init__(self, parent=None, depth=-1, value=None, init_score=1e6):
        # 这里的prior_p的先验设置功能还没写，动态剪肢也没有写
        # 层级指针
        self.depth = depth
        self.Parent = parent
        self.children = {}
        self.value = value
        self.min_score = init_score
        self.mean_score = init_score  # min
        self.n_visits = 0

    def expand(self, values, policy_value):
        """:cvar
        扩展时以子节点的变量值作为其value，在父结点中表示为key
        由于初始时会对子节点进行一次仿真，所以子节点的初值也会同时进行设置
        """
        self.children[values] = TreeNode(
            parent=self, depth=self.depth + 1, value=values)
        self.children[values].update_recursive(policy_value)

    def select(self, c_values, values_set=set(), thre=0.6, is_output=False):
        """:var
        选择结点
        """
        if is_output:
            children_score = np.array(
                list(map(lambda node: node.min_score, self.children.values())))
            selected_node = np.argmin(children_score)
            return list(self.children.items())[selected_node]
        else:
            children_value_set = set(self.children.keys())
            # 对于未完全展开的节点进行的处理
            if children_value_set != values_set:
                unfind_value_set = values_set - children_value_set
                temp_rate = np.random.rand()
                if temp_rate <= thre:  # 值越大广度越优先
                    choiced_value = np.random.choice(list(unfind_value_set))
                    return choiced_value, None

            children_score = np.array(
                list(map(lambda node: node.get_value(c_values), self.children.values())))  # 获得每个子节点的得分
            children_prob = self.softmax(children_score).cumsum()  # 计算其对应的概率

            choose_score = np.random.rand()
            selected_node = np.argmax(children_prob >= choose_score)
            children_value_list = list(self.children.keys())
            choiced_value = children_value_list[selected_node]
        return choiced_value, self.children[choiced_value]

    def update(self, new_value):
        """:var
        只有当其为未访问过的叶节点时才会调用该方法
        a = a+(b-1)/(n+1)
        """
        if self.n_visits == 0:
            self.min_score = new_value[0]
            self.mean_score = new_value[1]
        else:
            self.min_score = min(self.min_score, new_value[0])
            self.mean_score += (new_value[1] - self.mean_score) / self.n_visits
        self.n_visits += 1

    def update_recursive(self, new_value):
        """:var
        递归的更新每个结点的值
        """
        if not self.is_root():
            self.Parent.update_recursive(new_value)
        self.update(new_value)

    def get_value(self, c_values, eps=1e-8):
        """:var
        这里需要设计一个结点打分器
        需要考虑到结点的访问次数和当前得分
        先参考棋盘游戏中蒙特卡洛树的结构
        value = u + avg_score
        """
        # u = c_values * np.sqrt(np.log(self.Parent.n_visits + eps) / (self.n_visits + eps)) #先用UCB的方法
        u = c_values*np.sqrt(self.Parent.n_visits) / \
            (self.n_visits + eps)  # 先用UCB的方法
        return - (0.2*self.mean_score+self.min_score) + u  # 找最小值

    def is_leaf(self):
        return self.children == {}

    def is_root(self):
        return self.Parent is None
